package com.lingo.mgr.listener;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.lingo.mgr.domain.ChatSessionRecord;
import com.lingo.mgr.domain.vo.SseVo;
import com.lingo.mgr.enums.ChatMessageStatusEnum;
import com.lingo.mgr.enums.ChatgptError;
import com.lingo.mgr.mapper.ChatSessionRecordMapper;
import com.unfbx.chatgpt.entity.chat.ChatChoice;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.Message;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * sse
 */
@Slf4j
public class SseEmitterEventSessionListener extends EventSourceListener {

    /**
     * 数据接收结束
     */
    private static final String DONE = "[DONE]";

    /**
     * 数据接收完成
     */
    private static final String STOP = "stop";


    private final SseEmitter sseEmitter;

    private final ChatSessionRecordMapper chatSessionRecordMapper;

    private final StringBuilder stringBuilder;
    private final StringBuilder seeStringBuilder;



    private final Long dataId;

    public SseEmitterEventSessionListener(SseEmitter sseEmitter, Long dataId, ChatSessionRecordMapper chatSessionRecordMapper) {
        this.sseEmitter = sseEmitter;
        stringBuilder = new StringBuilder();
        seeStringBuilder = new StringBuilder();
        this.dataId = dataId;
        this.chatSessionRecordMapper = chatSessionRecordMapper;
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        log.info("ChatGpt建立sse连接...");
    }

    /**
     * 数据处理
     *
     * @param eventSource eventSource
     * @param id          事件id
     * @param type        类型
     * @param data        数据
     */
    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {
        log.info("ChatGpt返回数据：{}", data);
        if (DONE.equals(data)) {
            log.info("ChatGpt返回数据结束了");
            sseEmitter.complete();
        } else {
            try {
                // 解析data
                ObjectMapper mapper = new ObjectMapper();
                ChatCompletionResponse chatCompletionResponse = mapper.readValue(data, ChatCompletionResponse.class);
                if (null != chatCompletionResponse) {
                    List<ChatChoice> choices = chatCompletionResponse.getChoices();
                    if (CollectionUtils.isNotEmpty(choices)) {
                        ChatChoice choice = choices.get(0);
                        if (null != choice) {
                            String finishReason = choice.getFinishReason();
                            if (StringUtils.isNotEmpty(finishReason) && STOP.equals(finishReason)) {
                                log.info("ChatGpt返回stop, 数据结束了");
                                sseEmitter.complete();
                                return;
                            }
                            Message delta = choice.getDelta();
                            if (null != delta) {
                                String answerContent = delta.getContent();
                                if (null != answerContent) {
                                    seeStringBuilder.append(answerContent);
                                    SseVo sseVo = new SseVo(chatCompletionResponse.getId(), seeStringBuilder.toString());
                                    String answerStr = JSONUtil.toJsonStr(sseVo);
                                    log.info("sseEmitter send body is [{}]", answerStr);
                                    sseEmitter.send(SseEmitter.event().id(chatCompletionResponse.getId()).data(answerStr).reconnectTime(10000));
                                    TimeUnit.MILLISECONDS.sleep(RandomUtil.randomInt(10, 80));
                                    stringBuilder.append(answerContent);
                                }
                            }
                        }
                    }
                }
            } catch (IOException e) {
                log.error("sse发送成功消息给前端IO异常", e);
                eventSource.cancel();
            } catch (Exception e) {
                log.error("sse发送成功消息给前端未知异常", e);
                eventSource.cancel();
            }
        }
    }

    /**
     * 关闭
     *
     * @param eventSource eventSource
     */
    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        String answer = StringUtils.isEmpty(stringBuilder.toString()) ? "" : stringBuilder.toString();
        Message assistant = Message.builder().role(Message.Role.ASSISTANT).content(answer).build();
        log.info("ChatGpt返回数据:{}",JSONObject.toJSONString(assistant));
        chatSessionRecordMapper.update(ChatSessionRecord.builder()
                .answerResult(ChatMessageStatusEnum.SUCCESS.getCode())
                .messageAnswer(JSONObject.toJSONString(assistant)).build(),
                new LambdaUpdateWrapper<ChatSessionRecord>().eq(ChatSessionRecord::getSessionRecordId,dataId));
        log.info("ChatGpt关闭sse连接...");
        sseEmitter.complete();
    }

    /**
     * 失败
     *
     * @param eventSource eventSource
     * @param t           异常
     * @param response    响应
     */
    @Override
    public void onFailure(@NotNull EventSource eventSource, Throwable t, Response response) {
        String answer = StringUtils.isEmpty(stringBuilder.toString()) ? "" : stringBuilder.toString();
        Message assistant = Message.builder().role(Message.Role.ASSISTANT).content(answer).build();
        chatSessionRecordMapper.update(ChatSessionRecord.builder()
                .answerResult(ChatMessageStatusEnum.FAIL.getCode()).messageAnswer(JSONObject.toJSONString(assistant)).build(),
            new LambdaUpdateWrapper<ChatSessionRecord>().eq(ChatSessionRecord::getSessionRecordId,dataId));
        try {
            // 响应为空
            if (Objects.isNull(response)) {
                SseVo sseVo = new SseVo(IdUtil.simpleUUID(), ChatgptError.OPENAI_REQUEST_ERROR);
                log.error("ChatGPT异常且响应为空", t);
                sseEmitter.send(SseEmitter.event().id(sseVo.getId()).data(JSONUtil.toJsonStr(sseVo)).reconnectTime(10000));
                eventSource.cancel();
                sseEmitter.complete();
                return;
            }
            ResponseBody body = response.body();
            if (body != null) {
                String strBody = body.string();
                log.error("ChatGPT异常, body: {}", strBody, t);
                SseVo sseVo = new SseVo(IdUtil.simpleUUID(), strBody);
                sseEmitter.send(SseEmitter.event().id(sseVo.getId()).data(JSONUtil.toJsonStr(sseVo)).reconnectTime(10000));
            } else {
                log.error("ChatGPT异常且body为空, response: {}", response, t);
                SseVo sseVo = new SseVo(IdUtil.simpleUUID(), ChatgptError.OPENAI_REQUEST_ERROR);
                sseEmitter.send(SseEmitter.event().id(sseVo.getId()).data(JSONUtil.toJsonStr(sseVo)).reconnectTime(10000));
            }
            eventSource.cancel();
            sseEmitter.complete();
        } catch (IOException e) {
            log.error("sse发送失败消息给前端IO异常", e);
        } catch (Exception e) {
            log.error("sse发送失败消息给前端未知异常", e);
        }
    }
}
